# -*- coding: utf-8 -*-

"""
@Datetime: 2019/1/31
@Author: Zhang Yafei
"""
import json
import pymongo
import pandas as pd


class MongoPipeline(object):
    """
    mongodb:
        save(self, data, collection):                    将数据保存到数据库
        read(self, data):                                读取数据库中指定表格
        insert(self, table, dict_data)：                 插入数据
        delete(self, table, condition)：                 删除指定数据
        update(self, table, condition, new_dict_data)：  更新指定数据
        dbFind(self, table, condition=None):             按条件查找
        findAll(self, table)：                           查找全部
        close(self)：                                    关闭连接
    """

    def __init__(self, mongo_db, mongo_uri='localhost'):
        self.mongo_uri = mongo_uri
        self.mongo_db = mongo_db
        self.client = pymongo.MongoClient(self.mongo_uri)
        self.db = self.client[self.mongo_db]

    def close(self):
        """
        关闭连接
        :return:
        """
        self.client.close()

    def save(self, data, collection):
        """
        将数据保存到数据库表
        :param data:
        :param collection:
        :return: None
        """
        self.collection = self.db[collection]
        try:
            if self.collection.insert(json.loads(data.T.to_json()).values()):
                print('mongodb insert {} sucess.'.format(collection))
                return
        except Exception as e:
            print('insert error:', e)
            import traceback
            traceback.print_exc(e)

    def read(self, table):
        """
        读取数据库中的数据
        :param table:
        :return: dataframe
        """
        try:
            # 连接数据库
            table = self.db[table]
            # 读取数据
            data = pd.DataFrame(list(table.find()))
            return data
        except Exception as e:
            import traceback
            traceback.print_exc(e)

    def insert(self, table, dict_data):
        """
        插入
        :param table:
        :param dict_data:
        :return: None
        """
        try:
            self.db[table].insert(dict_data)
            print("插入成功")
        except Exception as e:
            print(e)

    def update(self,table, condition, new_dict_data):
        """
        更新
        :param table:
        :param dict_data:
        :param new_dict_data:
        :return: None
        """
        try:
            self.db[table].update(condition, new_dict_data)
            print("更新成功")
        except Exception as e:
            print(e)

    def delete(self, table, condition, one=False):
        """
        按条件删除
        :param table:
        :param dict_data:
        :return: None
        """
        try:
            if one:
                self.db[table].delete_one(condition)
                print('删除成功')
            else:
                result = self.db[table].delete_many(condition)
                print('删除成功')
                return result
        except Exception as e:
            print(e)

    def drop(self, table):
        """
        删除整个集合
        :param table:
        :return:
        """
        try:
            self.db[table].drop()
            print('删除{}成功'.format(table))
        except Exception as e:
            print(e)

    def dbFind(self, table, condition=None):
        """
        按条件查找
        :param table:
        :param dict_data:
        :return: generator dict
        """
        data = self.db[table].find(condition)
        for item in data:
            yield item

    def findAll(self, table):
        """
        查找全部
        :param table:
        :return: generator dict
        """
        for item in self.db[table].find():
            yield item


if __name__ == '__main__':
    mongo = MongoPipeline('flask')
    # data = mongo.read('label')
    # print(data.head())
    condition = {"药品ID": 509881}
    data = mongo.dbFind('label', condition)
    print(data)
    for i in data:
        print(i)
    # mongo.findAll()


